import time
import threading
import queue
from concurrent.futures import ThreadPoolExecutor
import traceback
from multiprocessing import cpu_count
from functools import partial
import os
from tqdm import tqdm

"""
多线程装饰器。

函数定义要求：
1 输入：
    datas: 第一个参数为待处理数据，为数组，或为一个文件（按行存储数据）

2 输出：
    res: 处理结果

装饰器使用要求：
1 res_path: 指定结果保存路径, 如果需要保存结果
"""


def add_queue(fn, que, res_handle_fn, *args, **kwargs):
    res = fn(*args, **kwargs)
    if res_handle_fn:
        res = res_handle_fn(res)
    if que:
        que.put(res)
    
    print(f'finish add queue')

# 消费者线程的目标函数，用于从队列中获取数据并写入文件
def write_from_queue(file_name, q):
    p_bar = tqdm()
    p_bar.set_description('write_from_queue:')
    while True:
        try:
            data = q.get()
            if data == "DONE":
                q.task_done()
                break
            with open(file_name, 'a', encoding='utf-8') as f:
                f.write(data + '\n')
            q.task_done()
            p_bar.update(1)

        except:
            print(f"write queue error: {traceback.format_exc()}")
    p_bar.close()

# 生产者任务的函数，用于将数据放入队列
def put_data_in_queue(data, q):
    q.put(data)

def test_fn(*args, **kwargs):
    print(f'test {args}, {kwargs}')


# 按行从文件读取数据
def read_line(path):
    with open(path, 'r', encoding='utf-8') as reader:
        line = reader.readline()
        while line:
            line = line.strip('\n')
            yield line

            line = reader.readline()


def multi_thread(logger=None, res_path=None, worker_nums=cpu_count(), res_handle_fn=lambda x:x):
  """
  The decorator for record time.
  The use: @fn_timer(logger) above the definition of method
  Args:
    logger: The logger for record log.

  Returns:
    (function):
  """
  def decorator(fn):
    """
    The decorator for record time.
    Args:
      fn(function): Function that is used by the decorator

    Returns:
      (funtion):
    """
    def running_multi_thread(*args, **kwargs):
        """
        The method multi thread.
        datas: args[0], 可为数组值，也可为文件路径（文件路径，按行读取值）
        Args:
            *args:
            **kwargs:

        Returns:

        """
        # 参数
        datas = args[0]
        if isinstance(datas, str) and os.path.isfile(datas):
            datas = read_line(datas)

        q = queue.Queue() if res_path else None
        if q:
            threading.Thread(target=write_from_queue, args=(res_path, q)).start()

        # 将运行结果放入 queue
        def add_queue(data, *args, **kwargs):
            res = fn(data, *args, **kwargs)
            if res_handle_fn:
                res = res_handle_fn(res)
            if q:
                q.put(res)
            
            # print(f'finish add queue')

        # 多线程开始
        with ThreadPoolExecutor(max_workers=worker_nums) as executor:
            for data in datas:
                executor.submit(add_queue, [data], *args[1:], **kwargs)
        
        # 所有生产者任务完成后，向队列中添加 "DONE" 消息
        if q:
            q.put("DONE")   
    
    return running_multi_thread

  return decorator



def multi_thread_no_return_value(logger=None, worker_nums=cpu_count(), res_handle_fn=lambda x:x):
  """
  The decorator for 无返回值的多线程.
  单线程方法直接转多线程。
  The use: @fn_timer(logger) above the definition of method
  Args:
    logger: The logger for record log.

  Returns:
    (function):
  """
  def decorator(fn):
    """
    The decorator for record time.
    Args:
      fn(function): Function that is used by the decorator

    Returns:
      (funtion):
    """
    def running_multi_thread(*args, **kwargs):
        """
        The method multi thread.
        datas: args[0], 可为数组值，也可为文件路径（文件路径，按行读取值）
        Args:
            *args:
            **kwargs:

        Returns:

        """
        # 参数
        datas = args[0]
        if isinstance(datas, str) and os.path.isfile(datas):
            datas = read_line(datas)

        # 多线程开始
        with ThreadPoolExecutor(max_workers=worker_nums) as executor:
            for data in datas:
                executor.submit(fn, [data], *args[1:], **kwargs)

    return running_multi_thread

  return decorator

@multi_thread(res_path='./log.txt')
def test2(data, a, *args, **kwargs):
    print(data, a, args, kwargs)
    return str(data)

@multi_thread_no_return_value()
def test3(data, a, *args, **kwargs):
    # print('test3',data, a, args, kwargs)
    print(f'test3: {data}, {a}, {args}, {kwargs}')
    return str(data)

def main():
    print('start')
    datas = list(range(10))
    # test2(datas, 4)
    test3(datas, 4)
    print('end')

if __name__ == '__main__':
    main()



